// #include "PCFG.h"
#include "PCFG_v2.h"
#include <chrono>
#include <fstream>
// #include "md5.h"
#include "SIMD_md5.h"
#include <iomanip>
#include <mpi.h>
using namespace std;
using namespace chrono;

// 编译指令
// mpicxx main.cpp train.cpp guessing.cpp SIMD_md5.cpp -o ../main  -fopenmp -O2

int main(int argc, char **argv)
{
  // 初始化MPI
  MPI_Init(&argc, &argv);

  int rank, size;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &size);
  
    double time_hash = 0;  // 用于MD5哈希的时间
    double time_guess = 0; // 哈希和猜测的总时长
    double time_train = 0; // 模型训练的总时长
    PriorityQueue q;
    double start_train = MPI_Wtime();
    q.m.train("/guessdata/Rockyou-singleLined-full.txt");
    q.m.order();
    double end_train = MPI_Wtime();
    time_train = end_train - start_train;

    q.init();
    cout << "here" << endl;
    int curr_num = 0;
    double start = MPI_Wtime();

    // cout << "Priority queue size: " << q.priority.size() << endl;
    // cout << "About to enter main loop" << endl;

    // 由于需要定期清空内存，我们在这里记录已生成的猜测总数
    int history = 0;
    // std::ofstream a("./files/results.txt");
    while (!q.priority.empty())
    {
      // cout << "Calling PopNext()" << endl;
      q.PopNext();
      // cout << "PopNext() completed" << endl;
      q.total_guesses = q.guesses.size();
      
      
      if(rank == 0){
        if (q.total_guesses - curr_num >= 100000)
        {
          cout << "Guesses generated: " << history + q.total_guesses << endl;
          curr_num = q.total_guesses;

          // 在此处更改实验生成的猜测上限
          int generate_n = 10000000;
          if (history + q.total_guesses > 10000000)
          {
            for (int dest = 1; dest < size; ++dest)
            {
              int terminate = -1;
              MPI_Send(&terminate, 1, MPI_INT, dest, 0, MPI_COMM_WORLD);
            }
            double end = MPI_Wtime();
            time_guess = end - start;
            cout << "Guess time:" << time_guess - time_hash << "seconds" << endl;
            cout << "Hash time:" << time_hash << "seconds" << endl;
            cout << "Train time:" << time_train << "seconds" << endl;
            break;
          }
          else{
            for (int dest = 1; dest < size; ++dest)
            {
              int terminate = 0;
              MPI_Send(&terminate, 1, MPI_INT, dest, 0, MPI_COMM_WORLD);
            }
          }
        }
        else{
          for (int dest = 1; dest < size; ++dest)
          {
            int terminate = 0;
            MPI_Send(&terminate, 1, MPI_INT, dest, 0, MPI_COMM_WORLD);
          }
        }
      }
      else{
        int terminate;
        MPI_Recv(&terminate, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        if(terminate == -1)
        {
          break;
        }
      }
      
      // 为了避免内存超限，我们在q.guesses中口令达到一定数目时，将其中的所有口令取出并且进行哈希
      // 然后，q.guesses将会被清空。为了有效记录已经生成的口令总数，维护一个history变量来进行记录
      if(rank == 0)
      {
        if (curr_num > 1000000)
        {
          double start_hash = MPI_Wtime();
          alignas(16) bit32 state[4][4];
          string pw[4];

// 使用 OpenMP 并行化外层循环
#pragma omp parallel for num_threads(7) shared(q) private(pw, state)
        for (int i = 0; i < q.guesses.size(); i += 4)
        {
          for (int j = i; j < q.guesses.size() && (j - i) < 4; j++)
          {
            pw[j - i] = q.guesses[j];
          }
          SIMD_MD5Hash(pw, state); // 假设 SIMD_MD5Hash 是线程安全的
        }

        // 在这里对哈希所需的总时长进行计算
        double end_hash = MPI_Wtime();
        time_hash += end_hash - start_hash;
        // auto duration = duration_cast<microseconds>(end_hash - start_hash);
        // time_hash += double(duration.count()) * microseconds::period::num / microseconds::period::den;

        // 记录已经生成的口令总数
        history += curr_num;
        curr_num = 0;
        q.guesses.clear();
      }
    }
  }
  // 完成MPI通信
  MPI_Finalize();
  return 0;
}